import cv2 as cv
import numpy as np

videoFilename = r'D:\opencv\opencv\sources\samples\data\vtest.avi'
# 角点检测参数
feature_params = dict(maxCorners = 100,
                      qualityLevel = 0.3,
                      minDistance = 7,
                      blockSize = 7 )
# lucas kanade光流法参数
lk_params = dict( winSize = (15, 15),
                  maxLevel = 2,
                  criteria = (cv.TERM_CRITERIA_EPS | cv.TERM_CRITERIA_COUNT, 10, 0.03))

cap = cv.VideoCapture(videoFilename)

#计算第一帧特征点
ret, prev = cap.read()
prevGray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)
p0 = cv.goodFeaturesToTrack(prevGray, mask=None, **feature_params)

while True:
    ret, frame =cap.read()
    if not ret: # 没读到当前帧，结束
        break

    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    # 计算光流
    p1, st, err = cv.calcOpticalFlowPyrLK(prevGray, gray, p0, None, **lk_params)

    # 选取好的跟踪点
    goodPoints = p1[st==1]
    goodPrevPoints = p0[st==1]
    res = frame.copy()
    drawColor = (0, 0, 255)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrevPoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv.line(res, (x0, y0), (x1, y1), drawColor)
        cv.circle(res, (x0, y0), 3, drawColor)

    # 更新上一帧
    prevGray =gray.copy()
    p0 = goodPoints.reshape(-1, 1, 2)

    # 显示计算结果图像
    cv.imshow('检测结果', res)

    key = cv.waitKey(30) # 每一帧间隔30ms
    if key == 27: # 按下ESC键， 退出
        break